{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7BdeKh8m8-tW"
      },
      "outputs": [],
      "source": [
        "# @title Install Code and Dependencies {form-width: \"25%\"}\n",
        "!pip install git+https://github.com/google-deepmind/tapnet.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 56,
          "status": "ok",
          "timestamp": 1739417525620,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 300
        },
        "id": "mCmDvfFvxnGB",
        "outputId": "76fa908f-1da1-442f-f3d2-59071052d274"
      },
      "outputs": [],
      "source": [
        "# @title Download Model {form-width: \"25%\"}\n",
        "\n",
        "%mkdir -p tapnet/checkpoints\n",
        "!wget -P tapnet/checkpoints --no-check-certificate https://storage.googleapis.com/dm-tapnet/tapnext/bootstapnext_ckpt.npz\n",
        "%ls tapnet/checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VADLwbmxaG29"
      },
      "outputs": [],
      "source": [
        "# @title Base Imports {form-width: \"10%\"}\n",
        "\n",
        "import io\n",
        "import os\n",
        "\n",
        "import cv2\n",
        "import einops\n",
        "import flax.linen as nn\n",
        "import jax\n",
        "import jax.nn as jnn\n",
        "import jax.numpy as jnp\n",
        "import matplotlib\n",
        "import mediapy as media\n",
        "import numpy as np\n",
        "from tapnet import evaluation_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HaswJZMq9B3c"
      },
      "outputs": [],
      "source": [
        "# @title Base Layers {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "class MlpBlock(nn.Module):\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    d = x.shape[-1]\n",
        "    x = nn.gelu(nn.Dense(4 * d)(x))\n",
        "    return nn.Dense(d)(x)\n",
        "\n",
        "\n",
        "class ViTBlock(nn.Module):\n",
        "  num_heads: int = 12\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    y = nn.LayerNorm()(x)\n",
        "    y = nn.MultiHeadDotProductAttention(num_heads=self.num_heads)(y, y)\n",
        "    x = x + y\n",
        "    y = nn.LayerNorm()(x)\n",
        "    y = MlpBlock()(y)\n",
        "    x = x + y\n",
        "    return x\n",
        "\n",
        "\n",
        "class Einsum(nn.Module):\n",
        "  width: int = 768\n",
        "\n",
        "  def setup(self):\n",
        "    self.w = self.param(\n",
        "        \"w\", nn.initializers.zeros_init(), (2, self.width, self.width * 4)\n",
        "    )\n",
        "    self.b = self.param(\n",
        "        \"b\", nn.initializers.zeros_init(), (2, 1, 1, self.width * 4)\n",
        "    )[:, 0]\n",
        "\n",
        "  def __call__(self, x):\n",
        "    return jnp.einsum(\"...d,cdD-\u003ec...D\", x, self.w) + self.b\n",
        "\n",
        "\n",
        "class RMSNorm(nn.Module):\n",
        "  width: int = 768\n",
        "\n",
        "  def setup(self):\n",
        "    self.scale = self.param(\"scale\", nn.initializers.zeros_init(), (self.width))\n",
        "\n",
        "  def __call__(self, x):\n",
        "    var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)\n",
        "    normed_x = x * jax.lax.rsqrt(var + 1e-6)\n",
        "    scale = jnp.expand_dims(self.scale, axis=range(len(x.shape) - 1))\n",
        "    return normed_x * (scale + 1)\n",
        "\n",
        "\n",
        "class Conv1D(nn.Module):\n",
        "  width: int = 768\n",
        "  kernel_size: int = 4\n",
        "\n",
        "  def setup(self):\n",
        "    self.w = self.param(\n",
        "        \"w\", nn.initializers.zeros_init(), (self.kernel_size, self.width)\n",
        "    )\n",
        "    self.b = self.param(\"b\", nn.initializers.zeros_init(), (self.width))\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    if state is None:\n",
        "      state = jnp.zeros(\n",
        "          (x.shape[0], self.kernel_size - 1, x.shape[1]), dtype=x.dtype\n",
        "      )\n",
        "    x = jnp.concatenate([state, x[:, None]], axis=1)  # shape: (b, k, c)\n",
        "    out = (x * self.w[None]).sum(axis=-2) + self.b[None]  # shape: (b, c)\n",
        "    state = x[:, 1 - self.kernel_size :]  # shape: (b, k - 1, c)\n",
        "    return out, state\n",
        "\n",
        "\n",
        "class BlockDiagonalLinear(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "\n",
        "  def setup(self):\n",
        "    width = self.width // self.num_heads\n",
        "    self.w = self.param(\n",
        "        \"w\", nn.initializers.zeros_init(), (self.num_heads, width, width)\n",
        "    )\n",
        "    self.b = self.param(\n",
        "        \"b\", nn.initializers.zeros_init(), (self.num_heads, width)\n",
        "    )\n",
        "\n",
        "  def __call__(self, x):\n",
        "    x = einops.rearrange(x, \"... (h i) -\u003e ... h i\", h=self.num_heads)\n",
        "    y = jnp.einsum(\"... h i, h i j -\u003e ... h j\", x, self.w) + self.b\n",
        "    return einops.rearrange(y, \"... h j -\u003e ... (h j)\", h=self.num_heads)\n",
        "\n",
        "\n",
        "class RGLRU(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "\n",
        "  def setup(self):\n",
        "    self.a_real_param = self.param(\n",
        "        \"a_param\", nn.initializers.zeros_init(), (self.width)\n",
        "    )\n",
        "    self.input_gate = BlockDiagonalLinear(\n",
        "        self.width, self.num_heads, name=\"input_gate\"\n",
        "    )\n",
        "    self.a_gate = BlockDiagonalLinear(self.width, self.num_heads, name=\"a_gate\")\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    gate_x = jnn.sigmoid(self.input_gate(x))\n",
        "    if state is None:\n",
        "      return gate_x * x  # No memory accumulation, return directly\n",
        "    else:\n",
        "      gate_a = jnn.sigmoid(self.a_gate(x))\n",
        "      log_a = -8.0 * gate_a * jnn.softplus(self.a_real_param)\n",
        "      a = jnp.exp(log_a)\n",
        "      scale_factor = jnp.sqrt(1 - jnp.exp(2 * log_a))  # Compute decay factor\n",
        "      return a * state + gate_x * x * scale_factor  # Memory update"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hICS3HPqcxU_"
      },
      "outputs": [],
      "source": [
        "# @title ViTSSM Blocks {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "class MLPBlock(nn.Module):\n",
        "  width: int = 768\n",
        "\n",
        "  def setup(self):\n",
        "    self.ffw_up = Einsum(self.width, name=\"ffw_up\")\n",
        "    self.ffw_down = nn.Dense(self.width, name=\"ffw_down\")\n",
        "\n",
        "  def __call__(self, x):\n",
        "    out = self.ffw_up(x)\n",
        "    return self.ffw_down(nn.gelu(out[0]) * out[1])\n",
        "\n",
        "\n",
        "class RecurrentBlock(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "\n",
        "  def setup(self) -\u003e None:\n",
        "    self.linear_y = nn.Dense(self.width, name=\"linear_y\")\n",
        "    self.linear_x = nn.Dense(self.width, name=\"linear_x\")\n",
        "    self.conv_1d = Conv1D(self.width, self.kernel_size, name=\"conv_1d\")\n",
        "    self.lru = RGLRU(self.width, self.num_heads, name=\"rg_lru\")\n",
        "    self.linear_out = nn.Dense(self.width, name=\"linear_out\")\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    y = jax.nn.gelu(self.linear_y(x))\n",
        "    x = self.linear_x(x)\n",
        "    x, conv1d_state = self.conv_1d(\n",
        "        x, None if state is None else state[\"conv1d_state\"]\n",
        "    )\n",
        "    rg_lru_state = self.lru(x, None if state is None else state[\"rg_lru_state\"])\n",
        "    x = self.linear_out(rg_lru_state * y)\n",
        "    return x, {\"rg_lru_state\": rg_lru_state, \"conv1d_state\": conv1d_state}\n",
        "\n",
        "\n",
        "class ResidualBlock(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "\n",
        "  def setup(self):\n",
        "    self.temporal_pre_norm = RMSNorm(self.width)\n",
        "    self.recurrent_block = RecurrentBlock(\n",
        "        self.width, self.num_heads, self.kernel_size, name=\"recurrent_block\"\n",
        "    )\n",
        "    self.channel_pre_norm = RMSNorm(self.width)\n",
        "    self.mlp = MLPBlock(self.width, name=\"mlp_block\")\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    y = self.temporal_pre_norm(x)\n",
        "    y, state = self.recurrent_block(y, state)\n",
        "    x = x + y\n",
        "    y = self.mlp(self.channel_pre_norm(x))\n",
        "    x = x + y\n",
        "    return x, state\n",
        "\n",
        "\n",
        "class ViTSSMBlock(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "\n",
        "  def setup(self):\n",
        "    self.ssm_block = ResidualBlock(self.width, self.num_heads, self.kernel_size)\n",
        "    self.vit_block = ViTBlock(self.num_heads)\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    b = x.shape[0]\n",
        "    x = einops.rearrange(x, \"b n c -\u003e (b n) c\")\n",
        "    x, state = self.ssm_block(x, state)\n",
        "    x = einops.rearrange(x, \"(b n) c -\u003e b n c\", b=b)\n",
        "    x = self.vit_block(x)\n",
        "    return x, state\n",
        "\n",
        "\n",
        "class ViTSSMBackbone(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "  num_blocks: int = 12\n",
        "\n",
        "  def setup(self):\n",
        "    self.blocks = [\n",
        "        ViTSSMBlock(\n",
        "            self.width,\n",
        "            self.num_heads,\n",
        "            self.kernel_size,\n",
        "            name=f\"encoderblock_{i}\",\n",
        "        )\n",
        "        for i in range(self.num_blocks)\n",
        "    ]\n",
        "    self.encoder_norm = nn.LayerNorm()\n",
        "\n",
        "  def __call__(self, x, state):\n",
        "    new_states = []\n",
        "    for i in range(self.num_blocks):\n",
        "      x, new_state = self.blocks[i](x, None if state is None else state[i])\n",
        "      new_states.append(new_state)\n",
        "    x = self.encoder_norm(x)\n",
        "    return x, new_states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BTUv4km7J514"
      },
      "outputs": [],
      "source": [
        "# @title TAPNext Model {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "def posemb_sincos_2d(h, w, width):\n",
        "  \"\"\"Compute 2D sine-cosine positional embeddings following MoCo v3 logic.\"\"\"\n",
        "  y, x = jnp.mgrid[0:h, 0:w]\n",
        "  freqs = jnp.linspace(0, 1, num=width // 4, endpoint=True)\n",
        "  inv_freq = 1.0 / (10_000**freqs)\n",
        "  y = jnp.einsum(\"h w, d -\u003e h w d\", y, inv_freq)\n",
        "  x = jnp.einsum(\"h w, d -\u003e h w d\", x, inv_freq)\n",
        "  pos_emb = jnp.concatenate(\n",
        "      [jnp.sin(x), jnp.cos(x), jnp.sin(y), jnp.cos(y)], axis=-1\n",
        "  )\n",
        "  return pos_emb\n",
        "\n",
        "\n",
        "class Backbone(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "  num_blocks: int = 12\n",
        "\n",
        "  def setup(self):\n",
        "    self.lin_proj = nn.Conv(\n",
        "        self.width,\n",
        "        (1, 8, 8),\n",
        "        strides=(1, 8, 8),\n",
        "        padding=\"VALID\",\n",
        "        name=\"embedding\",\n",
        "    )\n",
        "    self.mask_token = self.param(\n",
        "        \"mask_token\", nn.initializers.zeros_init(), (1, 1, 1, self.width)\n",
        "    )[:, 0]\n",
        "    self.unknown_token = self.param(\n",
        "        \"unknown_token\", nn.initializers.zeros_init(), (1, 1, self.width)\n",
        "    )\n",
        "    self.point_query_token = self.param(\n",
        "        \"point_query_token\", nn.initializers.zeros_init(), (1, 1, 1, self.width)\n",
        "    )[:, 0]\n",
        "    self.image_pos_emb = self.param(\n",
        "        \"pos_embedding\",\n",
        "        nn.initializers.zeros_init(),\n",
        "        (1, 256 // 8 * 256 // 8, self.width),\n",
        "    )\n",
        "    self.encoder = ViTSSMBackbone(\n",
        "        self.width,\n",
        "        self.num_heads,\n",
        "        self.kernel_size,\n",
        "        self.num_blocks,\n",
        "        name=\"Transformer\",\n",
        "    )\n",
        "\n",
        "  def __call__(self, frame, query_points, step, state):\n",
        "    x = self.lin_proj(frame)  # x: (b, h, w, c)\n",
        "    b, h, w, c = x.shape\n",
        "    query_points = jnp.concatenate(\n",
        "        [query_points[..., :1] - step, query_points[..., 1:]], axis=-1\n",
        "    )  # (b, q, 3)\n",
        "    posemb2d = posemb_sincos_2d(256, 256, self.width)  # posemb2d: (256, 256, c)\n",
        "\n",
        "    def interp(x, y):\n",
        "      return jax.scipy.ndimage.map_coordinates(\n",
        "          x, y.T - 0.5, order=1, mode=\"nearest\"\n",
        "      )\n",
        "\n",
        "    interp_fn = jax.vmap(interp, in_axes=(-1, None), out_axes=-1)\n",
        "    interp_fn = jax.vmap(interp_fn, in_axes=(None, 0), out_axes=0)\n",
        "    point_tokens = self.point_query_token + interp_fn(\n",
        "        posemb2d, query_points[..., 1:]\n",
        "    )  # (b, q, c)\n",
        "    # Query tokens\n",
        "    query_timesteps = query_points[..., 0:1].astype(jnp.int32)  # (b, q, 1)\n",
        "    query_tokens = jnp.where(\n",
        "        query_timesteps \u003e 0, self.unknown_token, self.mask_token\n",
        "    )  # (b, q, c)\n",
        "    query_tokens = jnp.where(\n",
        "        query_timesteps == 0, point_tokens, query_tokens\n",
        "    )  # (b, q, c)\n",
        "    # Image tokens\n",
        "    image_tokens = (\n",
        "        jnp.reshape(x, [b, h * w, c]) + self.image_pos_emb\n",
        "    )  # x: (b, h*w, c)\n",
        "\n",
        "    x = jnp.concatenate(\n",
        "        [image_tokens, query_tokens], axis=-2\n",
        "    )  # x: (b, h*w + q, c)\n",
        "    x, state = self.encoder(x, state)\n",
        "    _, q, _ = query_points.shape\n",
        "    x = x[:, -q:, :]  # x: (b, q, c)\n",
        "\n",
        "    return x, state\n",
        "\n",
        "\n",
        "class TAPNext(nn.Module):\n",
        "  width: int = 768\n",
        "  num_heads: int = 12\n",
        "  kernel_size: int = 4\n",
        "  num_blocks: int = 12\n",
        "\n",
        "  def setup(self):\n",
        "    self.backbone = Backbone(\n",
        "        self.width, self.num_heads, self.kernel_size, self.num_blocks\n",
        "    )\n",
        "    self.visible_head = nn.Sequential([\n",
        "        nn.Dense(256),\n",
        "        nn.LayerNorm(),\n",
        "        nn.gelu,\n",
        "        nn.Dense(256),\n",
        "        nn.LayerNorm(),\n",
        "        nn.gelu,\n",
        "        nn.Dense(1),\n",
        "    ])\n",
        "    self.coordinate_head = nn.Sequential([\n",
        "        nn.Dense(256),\n",
        "        nn.LayerNorm(),\n",
        "        nn.gelu,\n",
        "        nn.Dense(256),\n",
        "        nn.LayerNorm(),\n",
        "        nn.gelu,\n",
        "        nn.Dense(512),\n",
        "    ])\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, frame, query_points, step, state):\n",
        "    feat, state = self.backbone(frame, query_points, step, state)\n",
        "    track_logits = self.coordinate_head(feat)\n",
        "    visible_logits = self.visible_head(feat)\n",
        "\n",
        "    position_x, position_y = jnp.split(track_logits, 2, axis=-1)\n",
        "    position = jnp.stack([position_x, position_y], axis=-2)\n",
        "    index = jnp.arange(position.shape[-1])[None, None, None]\n",
        "    argmax = jnp.argmax(position, axis=-1, keepdims=True)\n",
        "    mask = jnp.abs(argmax - index) \u003c= 20\n",
        "    probs = jnn.softmax(position * 0.5, axis=-1) * mask\n",
        "    probs = probs / jnp.sum(probs, axis=-1, keepdims=True)\n",
        "    tracks = jnp.sum(probs * index, axis=-1) + 0.5\n",
        "    visible = (visible_logits \u003e 0).astype(jnp.float32)\n",
        "    return tracks, visible, state\n",
        "\n",
        "\n",
        "model = TAPNext()\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def forward(params, frame, query_points, step, state):\n",
        "  tracks, visible, state = model.apply(\n",
        "      {\"params\": params}, frame, query_points, step, state\n",
        "  )\n",
        "  return tracks, visible, state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YAvcpMST4g95"
      },
      "outputs": [],
      "source": [
        "# @title Load parameters {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "def npload(fname):\n",
        "  if os.path.exists(fname):\n",
        "    loaded = np.load(fname, allow_pickle=False)\n",
        "  else:\n",
        "    with open(fname, \"rb\") as f:\n",
        "      data = f.read()\n",
        "    loaded = np.load(io.BytesIO(data), allow_pickle=False)\n",
        "  if isinstance(loaded, np.ndarray):\n",
        "    return loaded\n",
        "  else:\n",
        "    return dict(loaded)\n",
        "\n",
        "\n",
        "def recover_tree(flat_dict):\n",
        "  tree = (\n",
        "      {}\n",
        "  )  # Initialize an empty dictionary to store the resulting tree structure\n",
        "  for (\n",
        "      k,\n",
        "      v,\n",
        "  ) in (\n",
        "      flat_dict.items()\n",
        "  ):  # Iterate over each key-value pair in the flat dictionary\n",
        "    parts = k.split(\n",
        "        \"/\"\n",
        "    )  # Split the key into parts using \"/\" as a delimiter to build the tree structure\n",
        "    node = tree  # Start at the root of the tree\n",
        "    for part in parts[\n",
        "        :-1\n",
        "    ]:  # Loop through each part of the key, except the last one\n",
        "      if (\n",
        "          part not in node\n",
        "      ):  # If the current part doesn't exist as a key in the node, create an empty dictionary for it\n",
        "        node[part] = {}\n",
        "      node = node[part]  # Move down the tree to the next level\n",
        "    node[parts[-1]] = v  # Set the value at the final part of the key\n",
        "  return tree  # Return the reconstructed tree\n",
        "\n",
        "\n",
        "ckpt_path = \"tapnet/checkpoints/bootstapnext_ckpt.npz\"\n",
        "params = recover_tree(npload(ckpt_path))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gyBJt05JWlRz"
      },
      "outputs": [],
      "source": [
        "# @title Visualization function {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "def plot_2d_tracks(\n",
        "    video,\n",
        "    points,\n",
        "    visibles,\n",
        "    infront_cameras=None,\n",
        "    tracks_leave_trace=16,\n",
        "    show_occ=False,\n",
        "):\n",
        "  \"\"\"Visualize 2D point trajectories.\"\"\"\n",
        "  num_frames, num_points = points.shape[:2]\n",
        "\n",
        "  # Precompute colormap for points\n",
        "  color_map = matplotlib.colormaps.get_cmap('hsv')\n",
        "  cmap_norm = matplotlib.colors.Normalize(vmin=0, vmax=num_points - 1)\n",
        "  point_colors = np.zeros((num_points, 3))\n",
        "  for i in range(num_points):\n",
        "    point_colors[i] = (np.array(color_map(cmap_norm(i)))[:3] * 255).astype(\n",
        "        np.uint8\n",
        "    )\n",
        "\n",
        "  if infront_cameras is None:\n",
        "    infront_cameras = np.ones_like(visibles).astype(bool)\n",
        "\n",
        "  frames = []\n",
        "  for t in range(num_frames):\n",
        "    frame = video[t].copy()\n",
        "\n",
        "    # Draw tracks on the frame\n",
        "    line_tracks = points[max(0, t - tracks_leave_trace) : t + 1]\n",
        "    line_visibles = visibles[max(0, t - tracks_leave_trace) : t + 1]\n",
        "    line_infront_cameras = infront_cameras[\n",
        "        max(0, t - tracks_leave_trace) : t + 1\n",
        "    ]\n",
        "    for s in range(line_tracks.shape[0] - 1):\n",
        "      img = frame.copy()\n",
        "\n",
        "      for i in range(num_points):\n",
        "        if line_visibles[s, i] and line_visibles[s + 1, i]:  # visible\n",
        "          x1, y1 = int(round(line_tracks[s, i, 0])), int(\n",
        "              round(line_tracks[s, i, 1])\n",
        "          )\n",
        "          x2, y2 = int(round(line_tracks[s + 1, i, 0])), int(\n",
        "              round(line_tracks[s + 1, i, 1])\n",
        "          )\n",
        "          cv2.line(frame, (x1, y1), (x2, y2), point_colors[i], 1, cv2.LINE_AA)\n",
        "        elif (\n",
        "            show_occ\n",
        "            and line_infront_cameras[s, i]\n",
        "            and line_infront_cameras[s + 1, i]\n",
        "        ):  # occluded\n",
        "          x1, y1 = int(round(line_tracks[s, i, 0])), int(\n",
        "              round(line_tracks[s, i, 1])\n",
        "          )\n",
        "          x2, y2 = int(round(line_tracks[s + 1, i, 0])), int(\n",
        "              round(line_tracks[s + 1, i, 1])\n",
        "          )\n",
        "          cv2.line(frame, (x1, y1), (x2, y2), point_colors[i], 1, cv2.LINE_AA)\n",
        "\n",
        "      alpha = (s + 1) / (line_tracks.shape[0] - 1)\n",
        "      frame = cv2.addWeighted(frame, alpha, img, 1 - alpha, 0)\n",
        "\n",
        "    # Draw end points on the frame\n",
        "    for i in range(num_points):\n",
        "      if visibles[t, i]:  # visible\n",
        "        x, y = int(round(points[t, i, 0])), int(round(points[t, i, 1]))\n",
        "        cv2.circle(frame, (x, y), 3, point_colors[i], -1, cv2.LINE_AA)\n",
        "      elif show_occ and infront_cameras[t, i]:  # occluded\n",
        "        x, y = int(round(points[t, i, 0])), int(round(points[t, i, 1]))\n",
        "        cv2.circle(frame, (x, y), 3, point_colors[i], 1, cv2.LINE_AA)\n",
        "\n",
        "    frames.append(frame)\n",
        "  frames = np.stack(frames)\n",
        "  return frames"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YNf98GICtJMa"
      },
      "outputs": [],
      "source": [
        "# @title Download TAPVid-DAVIS Dataset {form-width: \"25%\"}\n",
        "!wget --no-check-certificate https://storage.googleapis.com/dm-tapnet/tapvid_davis.zip\n",
        "!unzip tapvid_davis.zip"
      ]
    },
    {
      "metadata": {
        "id": "ecRl8zB2JOuO"
      },
      "cell_type": "code",
      "source": [
        "%%time\n",
        "# @title Qualitative results on dense tracks\n",
        "\n",
        "davis_dataset = evaluation_datasets.create_davis_dataset(\n",
        "    davis_points_path='./tapvid_davis/tapvid_davis.pkl',\n",
        "    query_mode='first',\n",
        "    full_resolution=False,\n",
        "    resolution=(256, 256),\n",
        ")\n",
        "\n",
        "for i, sample in enumerate(davis_dataset):\n",
        "  frames = sample['davis']['video']\n",
        "  ys, xs = np.meshgrid(np.linspace(8, 256, 32), np.linspace(8, 256, 32))\n",
        "  query_points = np.stack(\n",
        "      [np.zeros(len(xs.flatten())), xs.flatten(), ys.flatten()], axis=1\n",
        "  )[None]\n",
        "\n",
        "  tracks, visibles, state = [], [], None\n",
        "  for t in range(0, frames.shape[1]):\n",
        "    pred_tracks, pred_visible, state = forward(\n",
        "        params, frames[:, t], query_points, t, state\n",
        "    )\n",
        "    tracks.append(pred_tracks)\n",
        "    visibles.append(pred_visible)\n",
        "  tracks = np.stack(tracks, axis=2)[..., ::-1]\n",
        "  visibles = np.stack(visibles, axis=2).squeeze(-1)\n",
        "\n",
        "  frames = ((frames[0] + 1) / 2 * 255).astype(np.uint8)\n",
        "  video_viz = plot_2d_tracks(\n",
        "      frames, tracks[0].transpose(1, 0, 2), visibles[0].transpose(1, 0)\n",
        "  )\n",
        "  media.show_video(video_viz, fps=15)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "KcgAe_hmJOuO"
      },
      "cell_type": "code",
      "source": [
        "%%time\n",
        "# @title Qualitative results on sparse tracks\n",
        "\n",
        "davis_dataset = evaluation_datasets.create_davis_dataset(\n",
        "    davis_points_path='./tapvid_davis/tapvid_davis.pkl',\n",
        "    query_mode='first',\n",
        "    full_resolution=False,\n",
        "    resolution=(256, 256),\n",
        ")\n",
        "\n",
        "cnt = 0\n",
        "for i, sample in enumerate(davis_dataset):\n",
        "  batch = sample['davis']\n",
        "  tracks, visibles, state = [], [], None\n",
        "  for t in range(0, batch['video'].shape[1]):\n",
        "    pred_tracks, pred_visible, state = forward(\n",
        "        params, batch['video'][:, t], batch['query_points'], t, state\n",
        "    )\n",
        "    tracks.append(pred_tracks)\n",
        "    visibles.append(pred_visible)\n",
        "  tracks = np.stack(tracks, axis=2)[..., ::-1]\n",
        "  visibles = np.stack(visibles, axis=2).squeeze(-1)\n",
        "\n",
        "  frames = ((batch['video'][0] + 1) / 2 * 255).astype(np.uint8)\n",
        "  eye = np.eye(batch['video'].shape[1], dtype=np.int32)\n",
        "  query_frame_to_eval_frames = np.cumsum(eye, axis=1) - eye\n",
        "  query_frame = (batch['query_points'][0, :, 0]).astype(np.int32)\n",
        "  evaluation_mask = query_frame_to_eval_frames[query_frame]\n",
        "  visibles *= evaluation_mask[None]\n",
        "\n",
        "  video_viz = np.concatenate(\n",
        "      [\n",
        "          plot_2d_tracks(\n",
        "              frames, tracks[0].transpose(1, 0, 2), visibles[0].transpose(1, 0)\n",
        "          ),\n",
        "          plot_2d_tracks(\n",
        "              frames,\n",
        "              batch['target_points'][0].transpose(1, 0, 2),\n",
        "              ~batch['occluded'][0].transpose(1, 0),\n",
        "          ),\n",
        "      ],\n",
        "      axis=2,\n",
        "  )\n",
        "  media.show_video(video_viz, fps=15)\n",
        "\n",
        "  cnt += 1\n",
        "  if cnt \u003e 5:\n",
        "    break"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "yTEe1jK1JOuO"
      },
      "cell_type": "code",
      "source": [
        "%%time\n",
        "# @title Quantitative results\n",
        "\n",
        "davis_dataset = evaluation_datasets.create_davis_dataset(\n",
        "    davis_points_path='./tapvid_davis/tapvid_davis.pkl',\n",
        "    query_mode='first',\n",
        "    full_resolution=False,\n",
        "    resolution=(256, 256),\n",
        ")\n",
        "\n",
        "scores = []\n",
        "for i, sample in enumerate(davis_dataset):\n",
        "  batch = sample['davis']\n",
        "  print('video', i, jax.tree_util.tree_map(lambda x: x.shape, batch))\n",
        "  tracks, visibles, state = [], [], None\n",
        "  for t in range(0, batch['video'].shape[1]):\n",
        "    pred_tracks, pred_visible, state = forward(\n",
        "        params, batch['video'][:, t], batch['query_points'], t, state\n",
        "    )\n",
        "    tracks.append(pred_tracks)\n",
        "    visibles.append(pred_visible)\n",
        "  tracks = np.stack(tracks, axis=2)[..., ::-1]\n",
        "  visibles = np.stack(visibles, axis=2).squeeze(-1)\n",
        "\n",
        "  scalars = evaluation_datasets.compute_tapvid_metrics(\n",
        "      batch['query_points'],\n",
        "      batch['occluded'],\n",
        "      batch['target_points'],\n",
        "      np.logical_not(visibles),\n",
        "      tracks,\n",
        "      query_mode='first',\n",
        "  )\n",
        "  scalars = jax.tree.map(lambda x: np.array(np.sum(x, axis=0)), scalars)\n",
        "  scores.append(scalars)\n",
        "\n",
        "print(np.mean([scores[k]['average_jaccard'] for k in range(len(scores))]))\n",
        "print(np.mean([scores[k]['occlusion_accuracy'] for k in range(len(scores))]))\n",
        "print(\n",
        "    np.mean(\n",
        "        [scores[k]['average_pts_within_thresh'] for k in range(len(scores))]\n",
        "    )\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python/gpu:ml_notebook",
        "kind": "private"
      },
      "provenance": [
        {
          "file_id": "11QLNbzLDrzDRqO5YQNtzIyy8BA2OkjsF",
          "timestamp": 1737063264847
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
